import os
from datetime import datetime, timedelta
from typing import Dict, Any  # 添加 Any 到导入列表中
from typing import List
from typing import Tuple

import matplotlib.dates as mdates
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.preprocessing import MinMaxScaler
from sqlalchemy import create_engine
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau
from tensorflow.keras.layers import Input, LSTM, Dropout, Dense, RepeatVector, TimeDistributed
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers.legacy import Adam
from tensorflow.keras.utils import plot_model

plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False


class LSTMForecaster:
    def __init__(self, target_col='guilin_temp'):
        self.target_col = target_col.lower()
        self.db_table_name = 'weather_data'
        self.db_url = "mysql+pymysql://root:1026@localhost/guilin_temp?charset=utf8mb4"

        # 定义原始特征列（根据你的数据库实际列名调整）
        self.original_features = [
            'month', 'day', 'season', 'guilin_max', 'guilin_min', 'guilin_dewp',
            'guilin_prcp', 'guilin_slp', 'guilin_visib', 'guilin_wdsp', 'guilin_mxspd',
            'liuzhou_max', 'liuzhou_min', 'hezhou_max', 'hezhou_min'
        ]

        # 确保所有特征名都是小写
        self.original_features = [col.lower() for col in self.original_features]

        self.model = None
        self.scaler_features = MinMaxScaler(feature_range=(0, 1))
        self.scaler_targets = MinMaxScaler(feature_range=(0, 1))

        # 模型参数
        self.look_back = 180  # 使用90天历史数据
        self.forecast_horizon = 90  # 预测90天

        # 目录设置
        self.results_dir = os.path.join('app', 'static', 'results', 'lstm_forecast')
        os.makedirs(self.results_dir, exist_ok=True)
        self.model_dir = os.path.join('app', 'saved_models', 'LSTM')
        os.makedirs(self.model_dir, exist_ok=True)

    def load_and_preprocess_data(self) -> pd.DataFrame:
        """加载并预处理数据"""
        engine = create_engine(self.db_url)
        try:
            # 读取数据表
            df = pd.read_sql_table(self.db_table_name, engine)

            # 确保所有列名都是小写字符串
            df.columns = [str(col).lower() for col in df.columns]

            # 检查必要列是否存在
            required_columns = set(self.original_features + [self.target_col, 'record_date'])
            missing_cols = required_columns - set(df.columns)
            if missing_cols:
                raise ValueError(f"数据库中缺少必要的列: {missing_cols}")

            # 处理日期列
            df['record_date'] = pd.to_datetime(df['record_date'])
            df.set_index('record_date', inplace=True)

            # 数据清洗
            df = self._clean_data(df)
            # 特征工程
            df = self._feature_engineering(df)

            # 最终检查所有特征都是数值类型
            non_numeric = df.select_dtypes(exclude=['number', 'datetime']).columns
            if len(non_numeric) > 0:
                raise ValueError(f"以下列包含非数值数据: {list(non_numeric)}")

            return df
        except Exception as e:
            raise RuntimeError(f"数据加载失败: {str(e)}")
        finally:
            engine.dispose()

    def _clean_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """数据清洗"""
        # 处理目标列的缺失值
        if self.target_col in df.columns:
            df[self.target_col] = df[self.target_col].interpolate(method='time')
            # 去除异常值
            q_low = df[self.target_col].quantile(0.01)
            q_high = df[self.target_col].quantile(0.99)
            df[self.target_col] = df[self.target_col].clip(q_low, q_high)
        return df

    def _feature_engineering(self, df: pd.DataFrame) -> pd.DataFrame:
        """特征工程"""
        df = df.reset_index()
        # 确保所有列名都是字符串
        df.columns = [str(col) for col in df.columns]

        # 转换季节列为数值
        if 'season' in df.columns:
            season_map = {'spring': 0, 'summer': 1, 'fall': 2, 'winter': 3}
            df['season'] = df['season'].map(season_map).astype(float)

        # 时间特征
        df['dayofyear'] = df['record_date'].dt.dayofyear
        # 周期特征
        df['sin_year'] = np.sin(2 * np.pi * df['dayofyear'] / 365)
        df['cos_year'] = np.cos(2 * np.pi * df['dayofyear'] / 365)
        # 滞后特征
        for lag in [1, 7, 30, 90]:
            df[f'lag_{lag}'] = df[self.target_col].shift(lag)
        # 移动平均
        for window in [7, 30, 90]:
            df[f'ma_{window}'] = df[self.target_col].rolling(window).mean()

        # 确保所有特征都是数值类型
        for col in df.select_dtypes(include=['object']).columns:
            if col != 'record_date':  # 排除日期列
                try:
                    df[col] = pd.to_numeric(df[col])
                except ValueError:
                    df = df.drop(col, axis=1)  # 无法转换的列直接删除

        return df.dropna()

    def _create_dataset(self, data: pd.DataFrame) -> Tuple[np.ndarray, np.ndarray]:
        """创建时间序列数据集"""
        features = data.drop(columns=[self.target_col, 'record_date'])
        targets = data[[self.target_col]]

        # 标准化
        X = self.scaler_features.fit_transform(features)
        y = self.scaler_targets.fit_transform(targets)

        # 转换为3D格式 [samples, timesteps, features]
        X_seq, y_seq = [], []
        for i in range(len(X) - self.look_back - self.forecast_horizon + 1):
            X_seq.append(X[i:i + self.look_back])
            y_seq.append(y[i + self.look_back:i + self.look_back + self.forecast_horizon])

        return np.array(X_seq), np.array(y_seq)

    def train_test_split(self, test_size: float = 0.2) -> None:
        """划分训练测试集"""
        data = self.load_and_preprocess_data()
        self.dates = data['record_date'].values  # 保存日期数据

        X, y = self._create_dataset(data)
        split_idx = int(len(X) * (1 - test_size))

        self.X_train, self.X_test = X[:split_idx], X[split_idx:]
        self.y_train, self.y_test = y[:split_idx], y[split_idx:]
        self.dates_test = self.dates[split_idx + self.look_back:]  # 测试集对应日期

    def build_model(self) -> Model:
        """构建LSTM预测模型"""
        input_shape = (self.look_back, self.X_train.shape[2])

        inputs = Input(shape=input_shape)
        # 编码器
        x = LSTM(256, return_sequences=True)(inputs)
        x = Dropout(0.3)(x)
        x = LSTM(128, return_sequences=False)(x)
        # 解码器
        x = RepeatVector(self.forecast_horizon)(x)
        x = LSTM(128, return_sequences=True)(x)
        x = Dropout(0.3)(x)
        outputs = TimeDistributed(Dense(1))(x)

        self.model = Model(inputs, outputs)
        self.model.compile(optimizer=Adam(0.001), loss='mse')

        model_png_path = os.path.join(self.model_dir, 'lstm_model.png')

        plot_model(
            self.model,
            to_file=model_png_path,
            show_shapes=True,
            show_layer_names=True,
            rankdir='TB',
            dpi=300
        )
        print(f"模型结构图已保存至{model_png_path}")

        print(self.model.summary())

        return self.model

    def train(self, epochs: int = 100, batch_size: int = 32) -> dict:
        """训练模型"""
        callbacks = [
            EarlyStopping(patience=15, restore_best_weights=True),
            ModelCheckpoint(os.path.join(self.model_dir, 'new_best_model.h5'), save_best_only=True),
            ReduceLROnPlateau(factor=0.2, patience=5)
        ]

        history = self.model.fit(
            self.X_train, self.y_train,
            validation_split=0.2,
            epochs=epochs,
            batch_size=batch_size,
            callbacks=callbacks,
            verbose=1
        )
        self._plot_training_history(history)
        return history.history

    def _plot_training_history(self, history) -> None:
        """绘制训练过程曲线"""
        plt.figure(figsize=(12, 5))
        plt.plot(history.history['loss'], label='训练集损失')
        plt.plot(history.history['val_loss'], label='验证集损失')
        plt.title('模型训练过程', fontsize=14)
        plt.xlabel('训练轮次', fontsize=12)
        plt.ylabel('MSE损失值', fontsize=12)
        plt.legend(fontsize=12)
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.tight_layout()
        plt.savefig(os.path.join(self.results_dir, 'training_history.png'), dpi=300)
        plt.close()

    def load_model(self, weights_path: str) -> None:
        """
        加载预训练模型权重

        参数:
            weights_path: 模型权重文件路径(.h5文件)
        """
        if not os.path.exists(weights_path):
            raise FileNotFoundError(f"权重文件不存在: {weights_path}")

        # 先构建模型结构
        self.build_model()

        # 加载权重
        self.model.load_weights(weights_path)
        print(f"成功从 {weights_path} 加载模型权重")

    def evaluate(self) -> Dict[str, Any]:
        """评估模型性能并返回结构化结果

        返回:
            {
                'metrics': {
                    'MAE': float,
                    'RMSE': float,
                    'R2': float
                },
                'plot_path': str,  # 评估图表保存路径
                'dates': List[str],  # 日期列表(可选)
                'y_true': List[float],  # 实际值列表(可选)
                'y_pred': List[float]  # 预测值列表(可选)
            }
        """
        # 1. 进行预测
        y_pred = self.model.predict(self.X_test)

        # 2. 反标准化数据
        y_true = self.scaler_targets.inverse_transform(
            self.y_test[:, 0, :].reshape(-1, 1)).flatten()
        y_pred = self.scaler_targets.inverse_transform(
            y_pred[:, 0, :].reshape(-1, 1)).flatten()

        # 3. 计算评估指标
        metrics = {
            'MAE': round(mean_absolute_error(y_true, y_pred), 3),
            'RMSE': round(np.sqrt(mean_squared_error(y_true, y_pred)), 3),
            'R2': round(r2_score(y_true, y_pred), 3)
        }

        # 4. 生成评估图表
        plot_filename = f"lstm_evaluation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
        plot_path = os.path.join(self.results_dir, plot_filename)
        print(f'原始路径: {plot_path}')

        self._plot_evaluation(y_true, y_pred, plot_path)

        # 5. 构造返回结果
        dates = pd.to_datetime(self.dates_test[:len(y_true)]).astype(str).tolist()

        return {
            'metrics': metrics,
            'plot_path': plot_path,
            'dates': dates,  # 可选: 返回日期数据
            'y_true': y_true.tolist(),  # 可选: 返回实际值
            'y_pred': y_pred.tolist()  # 可选: 返回预测值
        }

    def _plot_evaluation(self, y_true: np.ndarray, y_pred: np.ndarray, save_path: str) -> None:
        """绘制预测结果对比图并保存到指定路径"""
        plt.figure(figsize=(14, 6))

        # 转换日期数据
        dates = pd.to_datetime(self.dates_test[:len(y_true)])

        # 智能日期标签
        ax = plt.gca()
        if len(dates) > 180:
            locator = mdates.MonthLocator(interval=3)
            formatter = mdates.DateFormatter('%Y-%m')
        else:
            locator = mdates.MonthLocator()
            formatter = mdates.DateFormatter('%m-%d')
        ax.xaxis.set_major_locator(locator)
        ax.xaxis.set_major_formatter(formatter)

        plt.plot(dates, y_true, label='实际值', linewidth=1, color='blue')
        plt.plot(dates, y_pred, label='预测值', linewidth=1, linestyle='--', color='red')

        plt.title('温度预测对比 (单步预测)', fontsize=14)
        plt.ylabel('温度 (°C)', fontsize=12)
        plt.legend(fontsize=12)
        plt.grid(True, linestyle=':', alpha=0.5)
        plt.gcf().autofmt_xdate()
        plt.tight_layout()

        # 保存到指定路径
        plt.savefig(save_path, dpi=300)
        plt.close()

    def forecast(self, days: int = 90) -> Dict[str, Any]:
        """预测未来指定天数的温度（适配接口版本）

        参数:
            days: 预测天数，默认为90天

        返回:
            {
                'status': 'success',
                'predictions': [
                    {'date': '2023-08-01', 'temperature': 28.5},
                    {'date': '2023-08-02', 'temperature': 28.7},
                    ...
                ],
                'plot_path': '/static/results/forecast_20230801.png',
                'statistics': {
                    'average': 26.8,
                    'max': 32.1,
                    'min': 22.5
                },
                'last_training_date': '2023-07-31'
            }
        """
        try:
            # 1. 加载并预处理数据
            data = self.load_and_preprocess_data()

            # 2. 获取最新特征数据
            last_features = data.drop(columns=[self.target_col, 'record_date']).values[-self.look_back:]

            # 3. 数据标准化和reshape
            X = self.scaler_features.transform(last_features)
            X = X.reshape(1, self.look_back, -1)

            # 4. 执行预测
            pred = self.model.predict(X)
            pred = self.scaler_targets.inverse_transform(pred[0].reshape(-1, 1)).flatten()

            # 5. 生成预测日期
            last_date = data['record_date'].iloc[-1]
            dates = [last_date + timedelta(days=i) for i in range(1, days + 1)]

            # 6. 生成预测结果字典
            predictions = [
                {
                    'date': d.strftime('%Y-%m-%d'),
                    'temperature': round(float(t), 1),
                    'day_of_week': d.strftime('%A'),
                    'day_of_year': d.timetuple().tm_yday
                }
                for d, t in zip(dates, pred[:days])
            ]

            # 7. 计算统计信息
            temps = [p['temperature'] for p in predictions]
            statistics = {
                'average': round(float(np.mean(temps)), 1),
                'max': round(float(np.max(temps)), 1),
                'min': round(float(np.min(temps)), 1)
            }

            # 8. 绘制预测图表
            plot_path = self._generate_forecast_plot(data, dates, pred[:days])

            return {
                'status': 'success',
                'predictions': predictions,
                'plot_path': plot_path,
                'statistics': statistics,
                'last_training_date': last_date.strftime('%Y-%m-%d'),
                'prediction_days': days,
                'generated_at': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            }

        except Exception as e:
            return {
                'status': 'error',
                'message': str(e)
            }

    def _generate_forecast_plot(self, history: pd.DataFrame, dates: List[datetime], pred: np.ndarray) -> str:
        """生成预测图表并返回文件路径"""
        plt.figure(figsize=(16, 6))

        # 绘制最近1年历史数据
        mask = history['record_date'] >= (history['record_date'].max() - pd.Timedelta(days=365))
        plt.plot(history.loc[mask, 'record_date'],
                 history.loc[mask, self.target_col],
                 label='历史数据',
                 linewidth=1.5,
                 color='#1f77b4')

        # 绘制预测数据
        plt.plot(dates, pred,
                 label=f'{len(pred)}天预测',
                 linestyle='--',
                 color='#ff7f0e',
                 linewidth=1.5)

        # 图表装饰
        ax = plt.gca()
        ax.xaxis.set_major_locator(mdates.MonthLocator(interval=2))
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
        plt.gcf().autofmt_xdate(rotation=45)

        plt.title(f'{self.target_col} {len(pred)}天温度预测', fontsize=14)
        plt.ylabel('温度 (°C)', fontsize=12)
        plt.legend(fontsize=12)
        plt.grid(True, linestyle=':', alpha=0.5)
        plt.tight_layout()

        # 确保结果目录存在
        os.makedirs(self.results_dir, exist_ok=True)

        # 生成唯一文件名
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        plot_filename = f'forecast_{len(pred)}days_{timestamp}.png'
        plot_path = os.path.join(self.results_dir, plot_filename)

        plt.savefig(plot_path, dpi=300, bbox_inches='tight')
        plt.close()

        # 返回相对于static目录的路径
        return os.path.join('results', 'lstm_forecast', plot_filename)
